package com.health.trans.writer;

import com.health.trans.channel.MemorySet;
import com.health.trans.common.util.GenericUtil;
import com.health.trans.kit.WriterConfig;
import com.health.trans.pojo.JobConfig;
import com.health.trans.service.AsyncDataService;
import com.health.trans.service.action.TransAction;
import lombok.Getter;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.batch.item.database.JpaItemWriter;
import org.springframework.dao.DataAccessResourceFailureException;
import org.springframework.orm.jpa.EntityManagerFactoryUtils;

import javax.persistence.EntityManager;
import javax.persistence.EntityManagerFactory;
import java.util.Collection;
import java.util.List;
import java.util.function.Consumer;

/**
 * @author nqlz
 */
public class CsJpaItemWriter<T> extends JpaItemWriter<T> {

    protected static final Log logger = LogFactory.getLog(CsJpaItemWriter.class);
    protected final JobConfig<T> jobConfig;
    protected final AsyncDataService asyncDataService;
    protected final WriterConfig<T> writerConfig;

    @Getter
    private final TransAction<? extends T, Object> transAction;
    protected EntityManagerFactory entityManagerFactory;
    private boolean usePersist = false;

    public CsJpaItemWriter(WriterConfig<T> writerConfig, AsyncDataService asyncDataService,EntityManagerFactory entityManagerFactory) {
        this.entityManagerFactory=entityManagerFactory;
        super.setEntityManagerFactory(entityManagerFactory);
        this.asyncDataService = asyncDataService;
        this.writerConfig = writerConfig;
        this.transAction = asyncDataService.getTransAction(writerConfig.getTableClass());
        this.jobConfig = writerConfig.getJobConfig();
        MemorySet memorySet = AsyncDataService.initChannel(writerConfig.getTableClass(), lineHandler);
        transAction.setMemorySet(memorySet);
    }

    @Override
    public void write(List<? extends T> items) {
        asyncDataService.syncData(items);
        jobConfig.log(items.size());
    }

    @Getter
    protected final Consumer<?> lineHandler = data -> csWrite(GenericUtil.get(data));

    public void csWrite(Collection<? extends T> items) {
        EntityManager entityManager = EntityManagerFactoryUtils.getTransactionalEntityManager(entityManagerFactory);
        if (entityManager == null) {
            throw new DataAccessResourceFailureException("Unable to obtain a transactional EntityManager");
        }
        doWrite(entityManager, items);
        entityManager.flush();
    }

    protected void doWrite(EntityManager entityManager, Collection<? extends T> items) {

        if (logger.isDebugEnabled()) {
            logger.debug("Writing to JPA with " + items.size() + " items.");
        }

        if (!items.isEmpty()) {
            long addedToContextCount = 0;
            for (T item : items) {
                if (!entityManager.contains(item)) {
                    if (usePersist) {
                        entityManager.persist(item);
                    } else {
                        entityManager.merge(item);
                    }
                    addedToContextCount++;
                }
            }
            if (logger.isDebugEnabled()) {
                logger.debug(addedToContextCount + " entities " + (usePersist ? " persisted." : "merged."));
                logger.debug((items.size() - addedToContextCount) + " entities found in persistence context.");
            }
        }

    }
}
